from flask import Flask, request, jsonify
from flask_cors import CORS  # 导入CORS
from pymongo import MongoClient
from celery_app import celery_app
from bson import ObjectId
from bson.errors import InvalidId
import datetime

app = Flask(__name__)
# 添加CORS支持，允许所有来源的跨域请求
CORS(app, resources={r"/*": {"origins": "*"}})

db = MongoClient("mongodb://root:cGbtfKNeSNwamdVl@127.0.0.1:27017/?authSource=admin").oj
subs = db.submissions
problems = db.problems

@app.post("/submit")
def submit():
    d = request.get_json(force=True)
    doc = {"problem_id": ObjectId(d["problem_id"]), "lang": d["lang"], "code": d["code"], "status": "Pending"}
    _id = subs.insert_one(doc).inserted_id
    celery_app.send_task("judge_code", args=[str(_id)])
    return jsonify({"submission_id": str(_id)})

@app.get("/result/<sid>")
def result(sid):
    doc = subs.find_one({"_id": ObjectId(sid)}, {'_id': 0})
    return jsonify(doc or {"error": "not found"})

@app.post("/rejudge/<sid>")
def rejudge(sid):
    subs.update_one({"_id": ObjectId(sid)}, {"$set": {"status": "Pending Rejudge"}})
    celery_app.send_task("judge_code", args=[str(sid)])
    return jsonify({"rejudge": "ok"})

@app.post("/problem/add")
def add_problem():
    try:
        data = request.get_json(force=True)
        
        # 检查必填字段
        required_fields = ["title", "difficulty", "time_limit", "memory_limit", 
                          "tags", "description", "input_description", 
                          "output_description", "samples"]
        for field in required_fields:
            if field not in data:
                return jsonify({"error": f"缺少必填字段: {field}"}), 400
                
        # 获取最新的问题编号
        latest_problem = problems.find_one(sort=[("problem_no", -1)])
        next_number = 1000 if not latest_problem or "problem_no" not in latest_problem else latest_problem["problem_no"] + 1
                
        # 准备问题文档
        problem_doc = {
            "problem_no": next_number,  # 添加题目编号
            "title": data["title"],
            "difficulty": data["difficulty"],
            "time_limit": data["time_limit"],
            "memory_limit": data["memory_limit"],
            "tags": data["tags"],
            "description": data["description"],
            "input_description": data["input_description"],
            "output_description": data["output_description"],
            "samples": data["samples"],
            "hint": data.get("hint", ""),
            "background":data.get("background",""),
            "source": data.get("source", ""),
            "creator": data.get("creator", "admin"),
            "create_time": datetime.datetime.now(),
            "is_deleted": False,
            "is_visible": True
        }
        # 插入数据库
        result = problems.insert_one(problem_doc)
        return jsonify({"success": True, "problem_id": str(result.inserted_id), "problem_no": next_number})
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.get("/problems")
def get_problems():
    try:
        # 获取分页参数
        page = int(request.args.get("page", 1))
        size = int(request.args.get("size", 10))
        
        # 获取用户ID（如果有的话）
        user_id = request.args.get("user_id", None)
        
        # 查询可见的非删除题目
        query = {"is_visible": True, "is_deleted": False}
        
        # 计算总数
        total = problems.count_documents(query)
        
        # 分页查询题目
        problem_list = problems.find(query).sort("problem_no", 1).skip((page - 1) * size).limit(size)
        
        result = []
        for problem in problem_list:
            # 计算通过率
            problem_id = problem["_id"]
            total_submissions = subs.count_documents({"problem_id": problem_id})
            passed_submissions = subs.count_documents({"problem_id": problem_id, "status": "Accepted"})
            
            pass_rate = "0%" if total_submissions == 0 else f"{(passed_submissions / total_submissions * 100):.1f}%"
            
            # 检查用户状态
            status = "unsolved"
            if user_id:
                user_submission = subs.find_one({"problem_id": problem_id, "user_id": user_id}, sort=[("_id", -1)])
                if user_submission:
                    status = user_submission["status"]
            
            # 构建响应数据
            problem_no = problem.get("problem_no", 0)
            problem_data = {
                "id": str(problem_id),  # 保留原始ID
                "problem_no": problem_no,  # 添加题目编号
                "title": problem['title'],  # 添加编号到标题
                "difficulty": problem["difficulty"],
                "tags": problem["tags"],
                "source": problem.get("source", ""),
                "passRate": pass_rate,
                "status": status
            }
            result.append(problem_data)
        
        return jsonify({
            "total": total,
            "page": page,
            "size": size,
            "problems": result
        })
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.get("/problem/<pid>")
def get_problem_detail(pid):
    try:
        # 尝试先按照编号查询
        if pid.isdigit() and int(pid) >= 1000:
            problem = problems.find_one({"problem_no": int(pid), "is_visible": True, "is_deleted": False})
        else:
            # 按照ID查询
            problem = problems.find_one({"_id": ObjectId(pid), "is_visible": True, "is_deleted": False})
        
        if not problem:
            return jsonify({"error": "题目不存在或已被删除"}), 404
        
        # 构建响应数据
        problem_no = problem.get("problem_no", 0)
        problem_data = {
            "id": str(problem["_id"]),
            "problem_no": problem_no,
            "title": f"题目 {problem_no}: {problem['title']}",
            "difficulty": problem["difficulty"],
            "time_limit": problem["time_limit"],
            "memory_limit": problem["memory_limit"],
            "tags": problem["tags"],
            "description": problem["description"],
            "background": problem.get("background", ""),
            "input_description": problem["input_description"],
            "output_description": problem["output_description"],
            "samples": problem["samples"],
            "hint": problem.get("hint", ""),
            "source": problem.get("source", ""),
            "create_time": problem["create_time"].isoformat() if "create_time" in problem else ""
        }
        
        return jsonify(problem_data)
    
    except InvalidId:
        return jsonify({"error": "无效的题目ID"}), 400
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)